### ################################################################################
### Purpose of this script:
### To make available a series of functions for chekcing model performance
### ################################################################################
here::i_am("newR/00_helpers.R")

get_mae3 <- function(yhat, y, type = "mean") {
    ybar <- apply(yhat, c(2, 3), mean)
    if (type == "mean") { 
        r <- mean(abs(ybar - y), na.rm = TRUE)
    } else {
        r <- median(abs(ybar - y), na.rm = TRUE)
    }
    return(r)    
}

get_calibration3 <- function(yhat, y, alpha = 0.05) {
    ci <- apply(yhat, c(2, 3), quantile, probs = c(alpha/2, 1-alpha/2), na.rm = TRUE)
    lo <- ci[1,,]
    hi <- ci[2,,]
    lo_le_actual <- lo <= y
    hi_ge_actual <- hi >= y
    coverage <- mean(lo_le_actual & hi_ge_actual)
    return(coverage)
}

get_pcp3 <- function(yhat, y) {
    ybar <- apply(yhat, c(2, 3), mean)
    predwin <- apply(ybar, 1, which.max)
    actual <- apply(y, 1, which.max)
    pcp <- mean(predwin == actual, na.rm = TRUE)
    pcp
}

get_brier3 <- function(yhat, y) {
    ## Get the winner for each iter/seat combo
    winner <- apply(yhat, c(1, 2), which.max)
    ## Get the probabilities by summarizing over the columns
    probs <- sapply(1:5, function(i) colMeans(winner == i))

    actual <- t(apply(y, 1, function(x) {
        winning_line <- max(x, na.rm = TRUE)
        as.numeric(x == winning_line)
    }))
    multi_class_brier <- (probs - actual) ^ 2
    sum(multi_class_brier, na.rm = TRUE) / nrow(probs)
}



get_mae <- function(mod, type = "mean") {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    
    fbar <- apply(f, c(2, 3), mean)
    y <- mod$data$y
 
}

get_mae2 <- function(mod, df, type = "mean") {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    
    fbar <- apply(f, c(2, 3), mean)
    candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                  "Cand_BE")
    candmat <- as.matrix(df[,candvars])

    ## Zero out
    fbar[which(candmat == 0)] <- 0
    ## Enforce sum to one constraint
    fbar <- fbar / rowSums(fbar)
    
    y <- mod$data$y
    if (type == "mean") { 
        r <- mean(abs(fbar - y))
    } else {
        r <- median(abs(fbar - y))
    }
    return(r)
}

get_calibration <- function(mod, df, alpha = 0.05) {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    ci <- apply(f, c(2, 3), quantile, probs = c(alpha/2, 1-alpha/2))
    actual <- mod$data$y

    ## Punch out instances where the 
    candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                       "Cand_BE")
    candmat <- as.matrix(df[,candvars])
    actual[which(candmat == 0)] <- 0
    
    lo <- ci[1,,]
    hi <- ci[2,,]
    lo_le_actual <- lo <= actual
    hi_ge_actual <- hi >= actual
    coverage <- mean(lo_le_actual & hi_ge_actual)
    return(coverage)
}

get_calibration2 <- function(mod, df, alpha = 0.05) {
    require(brms)
    f <- fitted(mod, summary = FALSE)

    candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                       "Cand_BE")
    candmat <- as.matrix(df[,candvars])

    for (i in 1:dim(f)[1]) {
        x <- f[i,,]
        x[which(candmat == 0, arr.ind = TRUE)] <- 0
        x <- x / rowSums(x)
        f[i,,] <- x
    }
    
    ci <- apply(f, c(2, 3), quantile, probs = c(alpha/2, 1-alpha/2))
    actual <- mod$data$y
    ### Replace eps values with true zeroes
    actual[which(actual == 1/40000)] <- 0

    lo_includes <- ci[1,,] <= actual
    hi_includes <- ci[2,,] >= actual
    coverage <- mean(lo_includes & hi_includes)
    return(coverage)
}

get_pcp <- function(mod, alpha = 0.05) {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    f_bar <- apply(f, c(2, 3), mean)
    predwin <- apply(f_bar, 1, which.max)
    actual <- apply(mod$data$y, 1, which.max)
    pcp <- mean(predwin == actual)
    pcp
}

get_pcp2 <- function(mod, df, alpha = 0.05) {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    f_bar <- apply(f, c(2, 3), mean)

    candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                       "Cand_BE")
    candmat <- as.matrix(df[,candvars])

    f_bar[which(candmat == 0)] <- 0
    f_bar <- f_bar / rowSums(f_bar)
    
    predwin <- apply(f_bar, 1, which.max)
    actual <- apply(mod$data$y, 1, which.max)
    pcp <- mean(predwin == actual)
    pcp
}

get_brier <- function(mod) {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    ## Get the winner for each iter/seat combo
    f_winner <- apply(f, c(1, 2), which.max)
    ## Get the probabilities by summarizing over the columns
    probs <- sapply(1:5, function(i) colMeans(f_winner == i))

    actual <- t(apply(mod$data$y, 1, function(x) {
        winning_line <- max(x, na.rm = TRUE)
        as.numeric(x == winning_line)
    }))
    multi_class_brier <- (probs - actual) ^ 2
    sum(multi_class_brier) / nrow(probs)
}


get_brier2 <- function(mod, df) {
    require(brms)
    f <- fitted(mod, summary = FALSE)
    
    candvars <- paste0(c("Con", "Lab", "Lib", "Nat", "Oth"),
                       "Cand_BE")
    candmat <- as.matrix(df[,candvars])

    for (i in 1:dim(f)[1]) {
        x <- pp[i,,]
        x[which(candmat == 0, arr.ind = TRUE)] <- 0
        x <- x / rowSums(x)
        f[i,,] <- x
    }

    ## Get the winner for each iter/seat combo
    f_winner <- apply(f, c(1, 2), which.max)
    ## Get the probabilities by summarizing over the columns
    probs <- sapply(1:5, function(i) colMeans(f_winner == i))

    actual <- t(apply(mod$data$y, 1, function(x) {
        winning_line <- max(x, na.rm = TRUE)
        as.numeric(x == winning_line)
    }))
    multi_class_brier <- (probs - actual) ^ 2
    sum(multi_class_brier) / nrow(probs)
}
